import re
import requests
import json
from urllib.parse import urlparse, parse_qs, unquote
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time
import os

class DouyinModel:
    def __init__(self):
        self.download_dir = 'downloads'
        if not os.path.exists(self.download_dir):
            os.makedirs(self.download_dir)

    def extract_url(self, text):
        """从文本中提取第一个URL链接"""
        url_pattern = r'https?://[^\s]+'
        urls = re.findall(url_pattern, text)
        return urls[0] if urls else None

    def download_video(self, video_url, progress_callback=None):
        """下载视频文件"""
        try:
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36',
                'Referer': 'https://www.douyin.com/'
            }
            
            response = requests.get(video_url, headers=headers, stream=True)
            response.raise_for_status()
            
            # 获取文件总大小
            total_size = int(response.headers.get('content-length', 0))
            block_size = 8192  # 8KB
            downloaded = 0
            
            filename = f"douyin_video_{int(time.time())}.mp4"
            filepath = os.path.join(self.download_dir, filename)
            
            with open(filepath, 'wb') as f:
                for chunk in response.iter_content(chunk_size=block_size):
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        if progress_callback and total_size > 0:
                            progress = int((downloaded / total_size) * 100)
                            progress_callback(progress)
                        
            return filepath
            
        except Exception as e:
            raise Exception(f"下载视频失败: {str(e)}")

    def get_video_url(self, share_url):
        """获取抖音分享链接中的视频真实下载链接"""
        try:
            chrome_options = Options()
            chrome_options.add_argument('--headless')
            chrome_options.add_argument('--disable-gpu')
            chrome_options.add_argument('--no-sandbox')
            chrome_options.add_argument('--disable-dev-shm-usage')
            chrome_options.add_argument('--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36')
            
            chromedriver_path = os.path.join(os.getcwd(), 'chromedriver', 'chromedriver-win64', 'chromedriver.exe')
            if not os.path.exists(chromedriver_path):
                raise FileNotFoundError(f"ChromeDriver not found at: {chromedriver_path}")
                
            service = Service(executable_path=chromedriver_path)
            driver = webdriver.Chrome(service=service, options=chrome_options)
            
            try:
                driver.get(share_url)
                time.sleep(5)
                page_source = driver.page_source
                
                patterns = [
                    r'https://v3-web\.douyinvod\.com/[^\s"\']+',
                    r'play_addr.*?url_list.*?\["(https://[^"]+)"',
                    r'video_url.*?"(https://[^"]+)"',
                    r'play_addr.*?url.*?"(https://[^"]+)"'
                ]
                
                for pattern in patterns:
                    matches = re.findall(pattern, page_source)
                    if matches:
                        video_url = matches[0]
                        video_url = video_url.replace('\\u0026', '&')
                        return video_url
                
                return None
                
            finally:
                driver.quit()
                
        except Exception as e:
            raise Exception(f"获取视频链接失败: {str(e)}") 